package cn.myeasyai.feature;

import java.io.IOException;
import java.io.InputStream;
import java.util.List;

import cn.myeasyai.util.MatrixUtil;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.ResNet50;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.springframework.stereotype.Service;

/**
 * 向量转换器
 * @author fushoujiang
 */
@Service
public class FeatureExtractor {
    private ComputationGraph model;

    public FeatureExtractor() throws IOException {
        try {
            ZooModel<ComputationGraph> zooModel = ResNet50.builder().build();
            this.model = (ComputationGraph)zooModel.initPretrained();
        } catch (Exception var2) {
            throw new IOException("Failed to initialize the pre-trained model: " + var2.getMessage(), var2);
        }
    }

    public INDArray extractFeatures(String imageFileUrl) throws IOException {
        NativeImageLoader loader = new NativeImageLoader(224L, 224L, 3L);
        INDArray image = loader.asMatrix(imageFileUrl);
        ImagePreProcessingScaler scaler = new ImagePreProcessingScaler(0.0D, 1.0D);
        scaler.transform(image);
        return this.model.outputSingle(new INDArray[]{image});
    }

    public List<Float> extractFeatures(InputStream imageFile) throws IOException {
        NativeImageLoader loader = new NativeImageLoader(224L, 224L, 3L);
        INDArray image = loader.asMatrix(imageFile);
        ImagePreProcessingScaler scaler = new ImagePreProcessingScaler(0.0D, 1.0D);
        scaler.transform(image);
        final INDArray indArray = this.model.outputSingle(new INDArray[]{image});
        return MatrixUtil.matrixToFloatList(indArray);
    }

}
